package com.katze.boot.plugins.redis.aspect;

import com.google.common.collect.ImmutableList;
import com.katze.boot.common.domain.LimitType;
import com.katze.boot.plugins.redis.RedisTemplateManager;
import com.katze.boot.plugins.redis.exception.LimitAccessException;
import com.katze.boot.plugins.redis.interfaces.Limit;
import com.katze.common.http.HttpUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.io.Serializable;
import java.lang.reflect.Method;
import java.util.Objects;
import java.util.concurrent.TimeUnit;


/**
 * 接口限流
 */
@Aspect
@Component
@ConditionalOnProperty(prefix = "spring.redis.limit", name = "enable", havingValue = "true")
public class LimitAspect implements InitializingBean {

    private RedisTemplate<String, Serializable> limitRedisTemplate;

    @Resource
    private RedisTemplateManager<Serializable> manager;

    @Value("${spring.redis.database:0}")
    private int database;

    @Value("${spring.redis.limit.database:}")
    private Integer limitDatabase;

    @Value("${spring.redis.limit.lua-mode: true}")
    private boolean luaMode;


    @Pointcut("@annotation(com.katze.boot.plugins.redis.interfaces.Limit)")
    public void pointcut() {
        // do nothing
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();

        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        LimitType limitType = limitAnnotation.limitType();
        String key;
        int limitPeriod = limitAnnotation.period();
        int limitCount = limitAnnotation.count();
        switch (limitType) {
            case IP:
                key = HttpUtils.getRemoteAddress(request);
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            default:
                key = method.getName().toUpperCase();
        }
        ImmutableList<String> keys = ImmutableList.of(String.join("_", limitAnnotation.prefix(), key));
        Long count;
        if (luaMode) {
            String luaScript = buildLuaScript();
            RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
            count = limitRedisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
        } else {
           count = limitRedisTemplate.boundValueOps(keys.get(0)).increment();
           if (count != null && count == 1) {
               limitRedisTemplate.expire(keys.get(0), limitPeriod, TimeUnit.SECONDS);
           }
        }
        if (count != null && count.intValue() <= limitCount) {
            return point.proceed();
        } else {
            throw new LimitAccessException(limitAnnotation.message());
        }
    }

    /**
     * 限流脚本
     * 调用的时候不超过阈值，则直接返回并执行计算器自加。
     *
     * @return lua脚本
     */
    private String buildLuaScript() {
        return "local c" +
                "\nc = redis.call('get',KEYS[1])" +
                "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
                "\nreturn c;" +
                "\nend" +
                "\nc = redis.call('incr',KEYS[1])" +
                "\nif tonumber(c) == 1 then" +
                "\nredis.call('expire',KEYS[1],ARGV[2])" +
                "\nend" +
                "\nreturn c;";
    }

    @Override
    public void afterPropertiesSet() {
        this.limitRedisTemplate = manager.select(limitDatabase == null? database:limitDatabase);
    }
}
